import xml.dom.minidom
import os

from .build import DatasetLoader


@DatasetLoader.register()
class LabelMeDataset:
    @classmethod
    def build(cls, dataset_config):
        dataset_root = dataset_config['dataset_root']

        dataset = []

        for image_dir in os.listdir(dataset_root):
            image_dir_path = os.path.join(dataset_root, image_dir)
            if not os.path.isdir(image_dir_path):
                continue

            for xml_filename in os.listdir(image_dir_path):
                if not xml_filename.endswith('.xml'):
                    continue

                xml_path = os.path.join(dataset_root, image_dir, xml_filename)
                collection = xml.dom.minidom.parse(xml_path).documentElement

                objects = collection.getElementsByTagName('object')

                image_filename = collection.getElementsByTagName('filename')[0].childNodes[0].data
                width = int(collection.getElementsByTagName('width')[0].childNodes[0].data)
                height = int(collection.getElementsByTagName('height')[0].childNodes[0].data)

                image_path = os.path.join(dataset_root, image_dir, image_filename)

                objects_list = []
                for object in objects:
                    label = object.getElementsByTagName('name')[0].childNodes[0].data
                    xmin = object.getElementsByTagName('bndbox')[0].getElementsByTagName('xmin')[0].childNodes[0].data
                    ymin = object.getElementsByTagName('bndbox')[0].getElementsByTagName('ymin')[0].childNodes[0].data
                    xmax = object.getElementsByTagName('bndbox')[0].getElementsByTagName('xmax')[0].childNodes[0].data
                    ymax = object.getElementsByTagName('bndbox')[0].getElementsByTagName('ymax')[0].childNodes[0].data


                    object_dict = {
                        'label': label,
                        'bounding_box': [float(xmin), float(ymin), float(xmax), float(ymax)],
                    }
                    objects_list.append(object_dict)

                # to uniform interface
                objects = []
                for object in objects_list:
                    object_dict = {
                        'box': object['bounding_box'],
                        'category': object['label']
                    }
                    objects.append(object_dict)

                record = {
                    'path': image_path,
                    'objects': objects,
                    'params': {
                        'height': height,
                        'width': width,
                    },
                    'transforms': [],
                }
                dataset.append(record)

        return dataset
